import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import OneHotEncoder
from sklearn.impute import SimpleImputer

# 🔧 MOD: 引入 XGBoost 和 LightGBM 模型
from xgboost import XGBRegressor
from lightgbm import LGBMRegressor

def ml_fill_missing_values(input_file, output_file, target_columns=None, model_type='random_forest'):
    """
    使用机器学习模型对缺失值指标进行预测填充。

    参数:
        input_file: 输入 CSV 文件路径
        output_file: 输出填充后 CSV 文件路径
        target_columns: 可选，需填充的目标列列表，默认为所有存在缺失值的列
        model_type: 可选，指定使用的模型类型，可选值：'random_forest', 'xgboost', 'lightgbm'
    """
    df = pd.read_csv(input_file, encoding='utf-8-sig')

    if target_columns is None:
        numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
        target_columns = [col for col in numeric_cols if df[col].isnull().any()]

    print(f"🎯 待填充目标列: {target_columns}")

    exclude_cols = ['序号', '国名En', '国名Ch'] + target_columns
    feature_cols = df.columns.difference(exclude_cols)

    encoder = OneHotEncoder(sparse=False, handle_unknown='ignore')
    encoded_countries = encoder.fit_transform(df[['国名Ch']])
    encoded_df = pd.DataFrame(encoded_countries, columns=encoder.get_feature_names_out(['国名Ch']))

    df['Year'] = df['Year'].astype(float)
    numeric_features = df[feature_cols].select_dtypes(include=[np.number])
    X_all = pd.concat([numeric_features.reset_index(drop=True), encoded_df.reset_index(drop=True)], axis=1)

    imputer = SimpleImputer(strategy='median')
    X_imputed = imputer.fit_transform(X_all)

    for target in target_columns:
        y = df[target]
        X_train = X_imputed[~y.isnull()]
        y_train = y.dropna()
        X_pred = X_imputed[y.isnull()]

        if len(y_train) > 0 and len(X_pred) > 0:
            # 🔧 MOD: 支持多模型选择
            if model_type == 'random_forest':
                model = RandomForestRegressor(n_estimators=100, random_state=42)
            elif model_type == 'xgboost':
                model = XGBRegressor(n_estimators=100, random_state=42, verbosity=0)
            elif model_type == 'lightgbm':
                model = LGBMRegressor(n_estimators=100, random_state=42)
            else:
                raise ValueError(f"不支持的模型类型: {model_type}")

            model.fit(X_train, y_train)
            y_pred = model.predict(X_pred)
            df.loc[y.isnull(), target] = y_pred
            print(f"✅ {target} 缺失值填充完成 ({len(y_pred)} 条记录) [使用模型: {model_type}]")
        else:
            print(f"⚠️ {target} 无足够数据训练或无缺失")

    df.to_csv(output_file, index=False, encoding='utf-8-sig')
    print(f"🎉 填充后的数据已保存：{output_file}")



#修改点已在代码中使用 # 🔧 MOD 注释标注。